import axios from 'axios'
import { logger } from './logger'

const COMFYUI_BASE_URL = process.env.COMFYUI_BASE_URL || 'http://localhost:8188'

// 执行 ComfyUI 工作流
export const executeComfyWorkflow = async (workflow: any, params: any) => {
    try {
        // 替换工作流中的参数
        const processedWorkflow = replaceWorkflowParams(workflow, params)

        // 发送工作流到 ComfyUI
        const response = await axios.post(`${COMFYUI_BASE_URL}/api/queue`, {
            prompt: processedWorkflow
        })

        // 等待工作流执行完成
        const result = await waitForWorkflowComplete(response.data.prompt_id)
        return result
    } catch (error) {
        logger.error('ComfyUI execution error:', error)
        throw new Error('工作流执行失败')
    }
}

// 替换工作流参数
const replaceWorkflowParams = (workflow: any, params: any) => {
    const processedWorkflow = JSON.parse(JSON.stringify(workflow))

    for (const nodeId in processedWorkflow.nodes) {
        const node = processedWorkflow.nodes[nodeId]
        if (node.inputs) {
            for (const key in node.inputs) {
                const paramId = `${nodeId}_${key}`
                if (params[paramId] !== undefined) {
                    node.inputs[key] = params[paramId]
                }
            }
        }
    }

    return processedWorkflow
}

// 等待工作流执行完成
const waitForWorkflowComplete = async (promptId: string) => {
    let retries = 0
    const maxRetries = 60 // 最多等待 5 分钟

    while (retries < maxRetries) {
        try {
            const response = await axios.get(`${COMFYUI_BASE_URL}/api/history/${promptId}`)
            if (response.data.status === 'completed') {
                return processWorkflowResult(response.data)
            }
            if (response.data.status === 'error') {
                throw new Error(response.data.error || '工作流执行失败')
            }
        } catch (error) {
            logger.error('Error checking workflow status:', error)
        }

        await new Promise(resolve => setTimeout(resolve, 5000)) // 每 5 秒检查一次
        retries++
    }

    throw new Error('工作流执行超时')
}

// 处理工作流结果
const processWorkflowResult = (result: any) => {
    // 根据结果类型返回不同格式
    if (result.output?.images?.length) {
        return {
            type: 'image',
            data: result.output.images[0]
        }
    }
    if (result.output?.text) {
        return {
            type: 'text',
            data: result.output.text
        }
    }
    return {
        type: 'json',
        data: result.output
    }
}